import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 加载数据集
trainset = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
traindata = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=0)

testset = torchvision.datasets.FashionMNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
testdata = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False, num_workers=0)

# 神经网络
class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, 5)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(10, 20, 5)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(20 * 4 * 4, 500)
        self.fc2 = nn.Linear(500, 200)
        self.fc3 = nn.Linear(200, 10)

    def forward(self, x):
        x = self.pool1(self.conv1(x))
        x = self.relu1(x)
        x = self.pool2(self.conv2(x))
        x = self.relu2(x)
        x = x.view(-1, 20 * 4 * 4)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x
net = Network().to(device)
net=Network()


# 优化器
#optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.8)
# 损失函数
loss = nn.CrossEntropyLoss()

epochs = 5
for epoch in range(epochs):
    for i, (x, y) in enumerate(traindata):
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        output = net(x)
        loss = torch.nn.functional.cross_entropy(output, y)
        loss.backward()
        optimizer.step()

        if (i+1) % 200 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                      .format(epoch + 1, epochs, i + 1, len(traindata),
                              loss.item()))

test_corrects = 0.0
test_num = 0
test_acc = 0.0
#只进行前向传播计算，不计算梯度，从而节省内存，加快运行速度
with torch.no_grad():
    for (X, Y) in testdata:
            output = net(X)  # 前向传播过程，输入为测试数据集，输出为对每个样本的预测
            _, pre_lab = torch.max(output.data, 1)  # 查找每一行中最大值对应的行标
            test_corrects += torch.sum(pre_lab == Y)  # 如果预测正确，则准确度val_corrects加1
            test_num += X.size(0)  # 当前用于训练的样本数量

test_acc = test_corrects.double().item() / test_num  # 计算在测试集上的分类准确率
print("test accuracy:", test_acc)